package org.shzhyun.model;

import java.nio.file.Paths;
import java.util.Map;

import lombok.extern.slf4j.Slf4j;
import org.shzhyun.tanslate.DenoiserTranslator;
import org.shzhyun.tanslate.TacotronSTFT;
import org.shzhyun.tanslate.TacotronTranslator;
import org.shzhyun.tanslate.WaveGlowTranslator;

import ai.djl.Device;
import ai.djl.inference.Predictor;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;


@Slf4j
@Component
public class ModelUtils {

	@Value("${tts.modulePath}")
	private String modulePath;

	private Criteria<NDList, NDArray> waveGlowEncoder;
	private Predictor<NDList, NDArray> waveGlowPredictor;
	private Predictor<NDArray, NDArray> stftPredictor;
	private Criteria<NDArray, NDArray> tacotronSTFT;
	private Criteria<NDList, Map> tacotronTranslator;
	private Predictor<NDList, Map> tacotronTranslatorPredictor;
	private Criteria<NDList, NDArray> denoiserEncoder;
	private Predictor<NDList, NDArray> denoiserPredictor;


//	@PostConstruct
//	public void InitModelUtils() {
//		log.info("===>ModelUtils=>2025-02-20=>10:47:14");
//		log.info("===>ModelUtils=>modulePath:{}",modulePath);
//        try {
//		} catch (Exception e) {
//			System.out.println("100");
//            e.printStackTrace();
//        }
//    }


	//通过网盘分享的文件：tts模型
	//链接: https://pan.baidu.com/s/1Hk9nh21rNqnycuB5g1Kxxg 提取码: mixr
	public NDArray get_mel_01(NDArray  wav) throws Exception{
		 System.out.println("===>get_mel=>ttsfilePath:"+modulePath);
		this.tacotronSTFT =
				Criteria.builder()
						.setTypes(NDArray.class, NDArray.class)
						.optTranslator(new TacotronSTFT())
						.optEngine("PyTorch")
						.optDevice(Device.cpu())
						.optModelPath(Paths.get(modulePath+"/tacotronSTFT.pt"))
						.build();
		stftPredictor = ModelZoo.loadModel(tacotronSTFT).newPredictor();
		System.out.println("10027");
		wav = wav.expandDims(0);  
		NDArray melspec = stftPredictor.predict(wav);
        //melspec = self.stft.mel_spectrogram(audio_norm)
        //melspec = torch.squeeze(melspec, 0)
		melspec = melspec.squeeze(0);
		//释放内存
		tacotronSTFT=null;
		stftPredictor=null;
		System.gc(); // 每次分配后建议GC

        return melspec;
	}
	public NDArray embed_02(NDArray frames_batch, Predictor<NDArray, NDArray> speakerEncoderPredictor) throws Exception{
		NDArray embed = speakerEncoderPredictor.predict(frames_batch);
       
       return embed;
	}
	public Map<String,NDArray> inference_chu_1(NDList frames_batch) throws Exception{
		tacotronTranslator =
				Criteria.builder()
						.setTypes(NDList.class, Map.class)
						.optTranslator(new TacotronTranslator())
						.optEngine("PyTorch")
						.optDevice(Device.cpu())
						.optModelPath(Paths.get(modulePath+"/tacotron2.pt"))
						.build();
		tacotronTranslatorPredictor = ModelZoo.loadModel(tacotronTranslator).newPredictor();
		System.out.println("10029");
		Map<String,NDArray> embed = tacotronTranslatorPredictor.predict(frames_batch);
       //System.out.println(embed);
		//释放内存
		tacotronTranslator = null;
		tacotronTranslatorPredictor = null;
		System.gc(); // 每次分配后建议GC

		return embed;
	}
	public NDArray denoiser_3(NDArray wav,NDManager manager) throws Exception{
		denoiserEncoder =
				Criteria.builder()
						.setTypes(NDList.class, NDArray.class)
						.optTranslator(new DenoiserTranslator())
						.optEngine("PyTorch")
						.optDevice(Device.cpu())
						.optModelPath(Paths.get(modulePath+"/denoiser.pt"))
						.build();
		denoiserPredictor = ModelZoo.loadModel(denoiserEncoder).newPredictor();
		System.out.println("10030");
		NDArray denoiser_strength = manager.create(1.0f);
		NDList input = new NDList(); 
		
		/*NDList dim = new NDList(); 
		dim.add(wav);
		NDArray adddim = NDArrays.stack(dim);*/
		
		input.add(wav); 
		input.add(denoiser_strength); 
		NDArray embed = denoiserPredictor.predict(input);
		//释放内存
		denoiserEncoder = null;
		denoiserPredictor = null;
		System.gc(); // 每次分配后建议GC

		return embed;
	}
	
	public NDArray generate_wave_2(NDArray mels_postnet,NDManager manager ) throws Exception{

		this.waveGlowEncoder =
				Criteria.builder()
						.setTypes(NDList.class, NDArray.class)
						.optTranslator(new WaveGlowTranslator())
						.optEngine("PyTorch")
						.optDevice(Device.cpu())
						.optModelPath(Paths.get(modulePath+"/waveGlow.pt"))
						.build();
		log.info("===>ModelUtils=>2025-02-20=>10:47:23");
		this.waveGlowPredictor = ModelZoo.loadModel(waveGlowEncoder).newPredictor();
		System.out.println("10026");
		System.out.println("26");
		NDArray sigma = manager.create(1.0);
		NDList input = new NDList();
		System.out.println("27");
		input.add(mels_postnet);
		input.add(sigma);
		System.out.println("28");
		NDArray embed = waveGlowPredictor.predict(input);

		//释放内存
		waveGlowEncoder = null;
		waveGlowPredictor = null;
		System.gc(); // 每次分配后建议GC
		System.out.println("29");
		return embed; 
	}
	
}
